#!/usr/bin/env python3
"""
将G3_query.json转换为Hugging Face datasets格式的数据处理器
"""

import json
import os
from typing import List, Dict, Any, Tuple
from datasets import Dataset, DatasetDict
import argparse
import pandas as pd



class G3DatasetConverter:
    """G3数据集转换器"""

    def __init__(self, json_file_path: str):
        """
        初始化转换器

        Args:
            json_file_path: G3_query.json文件路径
        """
        self.json_file_path = json_file_path
        self.data = []
        self.load_data()

    def load_data(self):
        """加载原始JSON数据"""
        print(f"正在加载数据文件: {self.json_file_path}")
        with open(self.json_file_path, 'r', encoding='utf-8') as f:
            self.data = json.load(f)
        print(f"成功加载 {len(self.data)} 条数据")

    def process_api_list(self, api_list: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """
        处理API列表，简化结构

        Args:
            api_list: 原始API列表

        Returns:
            处理后的API列表
        """
        processed_apis = []

        for api in api_list:
            # 处理参数，确保所有值都是一致的类型
            def process_params(params):
                if not params:
                    return []
                processed = []
                for param in params:
                    processed_param = {}
                    for key, value in param.items():
                        if key == 'default':
                            # 将default值统一转换为字符串以避免类型冲突
                            processed_param[key] = str(value) if value is not None else ''
                        else:
                            processed_param[key] = str(value) if value is not None else ''
                    processed.append(processed_param)
                return processed

            processed_api = {
                'category_name': str(api.get('category_name', '')),
                'tool_name': str(api.get('tool_name', '')),
                'api_name': str(api.get('api_name', '')),
                'api_description': str(api.get('api_description', '')),
                'method': str(api.get('method', 'GET')),
                'required_parameters': process_params(api.get('required_parameters', [])),
                'optional_parameters': process_params(api.get('optional_parameters', []))
            }
            processed_apis.append(processed_api)

        return processed_apis

    def format_for_hf_dataset(self) -> List[Dict[str, Any]]:
        """
        将数据格式化为Hugging Face Dataset格式

        Returns:
            格式化后的数据列表
        """
        formatted_data = []

        print("正在转换数据格式...")

        for i, item in enumerate(self.data):
            # 处理API列表
            processed_apis = self.process_api_list(item['api_list'])

            # 构建Hugging Face格式的数据
            formatted_item = {
                'query_id': int(item['query_id']),
                'query': str(item['query']),
                'task': str(item['query']),
                'api_list_ori': json.dumps(item['api_list'], ensure_ascii=False),  # 保留原始JSON字符串
                'api_list': processed_apis,        # 处理后的api_list
                'api_list_str': json.dumps(processed_apis, ensure_ascii=False),  # 字符串格式，便于模型处理
                'agent_name': 'toolbench_agent',
                'relevant_apis': item['relevant APIs'] or [],
                'relevant_apis_str': json.dumps(item['relevant APIs'] or [], ensure_ascii=False),
                'num_available_apis': len(processed_apis),
                'num_relevant_apis': len(item['relevant APIs'] or []),

                # 为API选择任务添加额外的格式化字段
                'api_names': [str(api['api_name']) for api in processed_apis],
                'api_descriptions': [str(api['api_description']) for api in processed_apis],
                'relevant_api_names': item['relevant APIs'] or [],  # 保持原有格式
            }

            formatted_data.append(formatted_item)

            if (i + 1) % 5000 == 0:
                print(f"已处理 {i + 1}/{len(self.data)} 条数据")
                print(formatted_item['agent_name'])

        print("数据格式转换完成!")
        return formatted_data

    def split_dataset(self, data: List[Dict[str, Any]],
                     train_ratio: float = 0.8,
                     val_ratio: float = 0.1,
                     test_ratio: float = 0.1,
                     random_seed: int = 42) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], List[Dict[str, Any]]]:
        """
        分割数据集

        Args:
            data: 原始数据
            train_ratio: 训练集比例
            val_ratio: 验证集比例
            test_ratio: 测试集比例
            random_seed: 随机种子

        Returns:
            训练集、验证集、测试集
        """
        import random

        assert abs(train_ratio + val_ratio + test_ratio - 1.0) < 1e-6, "比例之和必须为1"

        # 复制数据并打乱
        shuffled_data = data.copy()
        random.seed(random_seed)
        random.shuffle(shuffled_data)

        total_size = len(shuffled_data)
        train_size = int(total_size * train_ratio)
        val_size = int(total_size * val_ratio)

        train_data = shuffled_data[:train_size]
        val_data = shuffled_data[train_size:train_size + val_size]
        test_data = shuffled_data[train_size + val_size:]

        print(f"数据集分割完成:")
        print(f"  训练集: {len(train_data)} 条")
        print(f"  验证集: {len(val_data)} 条")
        print(f"  测试集: {len(test_data)} 条")

        return train_data, val_data, test_data

    def create_hf_dataset(self, train_ratio: float = 0.8,
                         val_ratio: float = 0.1,
                         test_ratio: float = 0.1,
                         save_path: str = None) -> DatasetDict:
        """
        创建Hugging Face Dataset

        Args:
            train_ratio: 训练集比例
            val_ratio: 验证集比例
            test_ratio: 测试集比例
            save_path: 保存路径

        Returns:
            Hugging Face DatasetDict
        """
        # 格式化数据
        formatted_data = self.format_for_hf_dataset()

        # 分割数据集
        train_data, val_data, test_data = self.split_dataset(
            formatted_data, train_ratio, val_ratio, test_ratio
        )

        # 创建Dataset对象
        train_dataset = Dataset.from_list(train_data)
        val_dataset = Dataset.from_list(val_data)
        test_dataset = Dataset.from_list(test_data)

        # 创建DatasetDict
        dataset_dict = DatasetDict({
            'train': train_dataset,
            'validation': val_dataset,
            'test': test_dataset
        })

        # 保存数据集
        if save_path:
            print(f"正在保存数据集到: {save_path}")
            dataset_dict.save_to_disk(save_path)
            print("数据集保存完成!")

        return dataset_dict

    def export_to_jsonl(self, data: List[Dict[str, Any]], output_path: str):
        """
        导出为JSONL格式

        Args:
            data: 要导出的数据
            output_path: 输出路径
        """
        print(f"正在导出JSONL文件到: {output_path}")

        with open(output_path, 'w', encoding='utf-8') as f:
            for item in data:
                f.write(json.dumps(item, ensure_ascii=False) + '\n')

        print("JSONL文件导出完成!")

    def export_to_parquet(self, data: List[Dict[str, Any]], output_path: str):
        """
        导出为Parquet格式

        Args:
            data: 要导出的数据
            output_path: 输出路径
        """
        print(f"正在导出Parquet文件到: {output_path}")

        # 转换为DataFrame
        df = pd.DataFrame(data)

        # 导出为Parquet格式
        df.to_parquet(output_path, engine='pyarrow', compression='snappy')

        print("Parquet文件导出完成!")
        print(f"文件大小: {os.path.getsize(output_path) / (1024*1024):.2f} MB")


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='转换G3数据集为Hugging Face格式')
    parser.add_argument('--input', type=str,
                       default='/mnt/nas/users/wangmin/data/mirrorapi/G3_query.json',
                       help='输入JSON文件路径')
    parser.add_argument('--output_dir', type=str,
                       default='./g3_hf_dataset',
                       help='输出目录')
    parser.add_argument('--train_ratio', type=float, default=0.8, help='训练集比例')
    parser.add_argument('--val_ratio', type=float, default=0.1, help='验证集比例')
    parser.add_argument('--test_ratio', type=float, default=0.1, help='测试集比例')
    parser.add_argument('--export_jsonl', action='store_true', help='同时导出JSONL格式')
    parser.add_argument('--export_parquet', action='store_true', help='同时导出Parquet格式')

    args = parser.parse_args()

    # 创建转换器
    converter = G3DatasetConverter(args.input)

    # 创建Hugging Face数据集
    dataset_dict = converter.create_hf_dataset(
        train_ratio=args.train_ratio,
        val_ratio=args.val_ratio,
        test_ratio=args.test_ratio,
        save_path=args.output_dir
    )

    # 显示数据集信息
    print("\n=== 数据集信息 ===")
    for split_name, dataset in dataset_dict.items():
        print(f"{split_name}: {len(dataset)} 条数据")

    print(f"\n数据样例 (第一条训练数据):")
    sample = dataset_dict['train'][0]
    for key, value in sample.items():
        if key in ['api_list', 'api_list_str', 'relevant_apis_str', 'api_list_ori']:
            print(f"{key}: {str(value)[:100]}...")
        else:
            print(f"{key}: {value}")

    # 获取格式化数据用于导出
    formatted_data = converter.format_for_hf_dataset()
    train_data, val_data, test_data = converter.split_dataset(formatted_data, 
                                                            train_ratio=args.train_ratio, 
                                                            val_ratio=args.val_ratio,
                                                            test_ratio=args.test_ratio)

    os.makedirs(args.output_dir, exist_ok=True)

    # 导出JSONL格式（可选）
    if args.export_jsonl:
        converter.export_to_jsonl(train_data, os.path.join(args.output_dir, 'train.jsonl'))
        converter.export_to_jsonl(val_data, os.path.join(args.output_dir, 'validation.jsonl'))
        converter.export_to_jsonl(test_data, os.path.join(args.output_dir, 'test.jsonl'))

    # 导出Parquet格式（可选）
    if args.export_parquet:
        converter.export_to_parquet(train_data, os.path.join(args.output_dir, 'train.parquet'))
        converter.export_to_parquet(val_data, os.path.join(args.output_dir, 'validation.parquet'))
        converter.export_to_parquet(test_data, os.path.join(args.output_dir, 'test.parquet'))


if __name__ == "__main__":
    main()
